import tensorflow as tf
from tensorflow.keras import layers
import tensorflow.keras.applications as models
import tensorflow_ranking as tfr
import tensorflow_recommenders as tfrs


IMAGENET_INPUT_H = 224
IMAGENET_INPUT_W = 224
IMAGENET_INPUT_D = 3

IMAGENET_IMAGE_SIZE = (IMAGENET_INPUT_H, IMAGENET_INPUT_W)
IMAGENET_SHAPE = (IMAGENET_INPUT_H, IMAGENET_INPUT_W, IMAGENET_INPUT_D)

NUM_CLASSES = 5


class EfficientNetB1RankingModel(tfrs.Model):
    def __init__(self, loss, trainable=False, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.loss_fn = loss

        self.base_model = models.EfficientNetB1(
            weights='imagenet', 
            include_top=False,
            input_shape=IMAGENET_SHAPE,
        )
        self.base_model.trainable = trainable

        self.flatten = tf.keras.layers.Flatten()

        self.dense_layers = tf.keras.Sequential([
            tf.keras.layers.Dense(512, activation='relu'),
            tf.keras.layers.Dropout(0.5),
            tf.keras.layers.Dense(256, activation='relu'),
            tf.keras.layers.Dropout(0.5),
            tf.keras.layers.Dense(128, activation='relu'),
            tf.keras.layers.Dropout(0.5),
            tf.keras.layers.Dense(64, activation='relu'),
            tf.keras.layers.Dropout(0.5),
        ])

        self.position_score_model = tf.keras.Sequential([
            tf.keras.layers.Dense(1)
        ])

        self.task = tfrs.tasks.Ranking(
            loss=self.loss_fn,
            metrics=[
                tfr.keras.metrics.NDCGMetric(name="ndcg_metric"),
                tfr.keras.metrics.MRRMetric(name="mrr_metric"),
                tfr.keras.metrics.OPAMetric(name="opa_metric"),
                # tf.keras.metrics.MeanSquaredError(name="mse_metric"), # non-representative
            ],
        )

    def get_config(self):
        config = super().get_config().copy()
        config.update({
            'loss_fn': self.loss_fn,
        })
        return config

    @classmethod
    def from_config(cls, config):
        return cls(**config)

    def call(self, inputs):
        flatten_inputs = tf.reshape(inputs, [-1, IMAGENET_INPUT_H, IMAGENET_INPUT_W, IMAGENET_INPUT_D])
        embeddings = self.base_model(flatten_inputs)
        flattened_embeddings = self.flatten(embeddings)
        dense_output = self.dense_layers(flattened_embeddings)
        scores = self.position_score_model(dense_output)
        gathered_scores = tf.reshape(scores, [-1, NUM_CLASSES])
        return gathered_scores
        

    def compute_loss(self, features, training=False):
        images, labels = features
        scores = self(images)
        #labels = tf.one_hot(labels, depth=NUM_CLASSES)
        
        labels_ = tf.cast(labels, dtype="float32")
        
        return self.task(
            labels=labels_,
            predictions=scores
        )
        

        
